# ==============================================================================

# TITLE: Replication R code for the paper entitled: "Electoral systems, partisan politics, and income redistribution: A critical quasi-experiment"

# Authors: Michał Pierzgalski (the corresponding author), Maciej A. Górecki
# Version: 0.9.9
# Date: February 12, 2023

# R version 4.2.2 (2022-10-31)
# RStudio version: 2022.12.0+353 for macOS
# Running under: macOS Ventura 13.1

# ==============================================================================

rm(list = ls())


# START ========================================================================

# Install packages -------------------------------------------------------------

devtools::install_github('liulch/bpCausal') # R package for A Bayesain Alternative to the Synthetic Control Method. Version: 0.0.1 (GitHub version)

# Load packages ----------------------------------------------------------------

library(tidyverse) # Ver. 1.3.2
library(panelView) # Ver. 1.1.11
library(bpCausal) # Ver. 0.0.1 (GitHub version)

library(ggpubr) # Ver. 0.5.0
library(coda) # Ver. 0.19-4

## Options ---------------------------------------------------------------------
options(scipen = 999)
options(digits = 3)



#####################
#### Preparation ####
#####################

# Load data sets !!! ===========================================================

# Download (from CPS Harvard Dataverse), and open in RStudio the following file (RStudio workspace): "raw_workspace_1.2.23.RData". The workspace contains the data sets (see (online) Appendix for the description of data sources).

load("raw_workspace_1.2.23.RData") # !!!


# DATA preparation (1) =========================================================

data_bpC <-
  filter(data.final.2,!(ccode %in% c("OECD", "IRL")) &
           (year <= 2019 &
              year >= 1960)) %>% select(
                year,
                ccode,
                vturn,
                gdppc_2015,
                inflation_cpi,
                inflation_cpds,
                gov_left1,
                gov_left2,
                gov_left3,
                gov_right1,
                gov_right2,
                gov_right3,
                gov_cent1,
                gov_cent2,
                gov_cent3,
                gov_sup,
                elderly,
                young,
                gov_party,
                womenpar,
                effpar_ele,
                effpar_leg,
                strike,
                debt_hist,
                interest,
                popul,
                ud_ipol,
                educexp_gov_ipol,
                oldage_pmp,
                realgdpgr,
                netu_ipol,
                ttl_labf,
                gov_new,
                lfirstp,
                structur,
                lfirstpb,
                lfirstpi,
                outlays,
                bic,
                pres,
                fed,
                eu,
                judrev,
                starts_with("ssocial"),
                starts_with("sconserv"),
                starts_with("sliberal"),
                sconlib,
                sconserv,
                max.sconserv,
                postfisc_gini,
                prefisc_gini,
                gini_disp,
                eneg,
                vp_rigidity,
                regime,
                enpres.2,
                instcons,
                rel_red,
                abs_red,
                gini_mkt,
                UESKGEN,
                log.UESKGEN,
                probit_gr1,
                logit_gr1,
                topone_inc,
                topOneP,
                health_pmp,
                govexp,
                comp_voting
              ) %>% mutate(
                rnetu = netu_ipol / ttl_labf,
                rnetu.2 = netu_ipol / popul,
                logit.sconserv = (sconserv / (100 - sconserv)),
                logit.gov_right1 = (gov_right1 / (100 - gov_right1))
              )


# DATA preparation (2) =========================================================

data_bpC.2 <- as.data.frame(data_bpC)

DATA <-
  data_bpC.2 %>% mutate(
    TreatNZL.93 = ifelse((year >= 1993 &
                            ccode == "NZL"), 1, 0),
    TreatNZL.92 = ifelse((year >= 1992 &
                            ccode == "NZL"), 1, 0),
    TreatNZL.91 = ifelse((year >= 1991 &
                            ccode == "NZL"), 1, 0),
    TreatNZL.87 = ifelse((year >= 1987 &
                            ccode == "NZL"), 1, 0),
    TreatNZL.97 = ifelse((year >= 1997 &
                            ccode == "NZL"), 1, 0),
    TreatNZL.96 = ifelse((year >= 1996 & ccode == "NZL"), 1, 0)
  )


# For placebo analyses =========================================================

DATA.plac <-
  data_bpC.2 %>% mutate(
    TreatGBR.93 = ifelse((year >= 1993 &
                            ccode == "GBR"), 1, 0),
    TreatGBR.94 = ifelse((year >= 1994 &
                            ccode == "GBR"), 1, 0),
    TreatGBR.84 = ifelse((year >= 1984 &
                            ccode == "GBR"), 1, 0),
    TreatGBR.87 = ifelse((year >= 1987 &
                            ccode == "GBR"), 1, 0),
    TreatGBR.97 = ifelse((year >= 1997 &
                            ccode == "GBR"), 1, 0),
    TreatGBR.96 = ifelse((year >= 1996 &
                            ccode == "GBR"), 1, 0),
    TreatJPN.97 = ifelse((year >= 1997 & ccode == "JPN"), 1, 0),
    TreatAUS.97 = ifelse((year >= 1997 & ccode == "AUS"), 1, 0),
    TreatFRA.97 = ifelse((year >= 1997 & ccode == "FRA"), 1, 0),
    TreatUSA.97 = ifelse((year >= 1997 & ccode == "USA"), 1, 0),
    TreatCAN.97 = ifelse((year >= 1997 & ccode == "CAN"), 1, 0)
  )

# ==============================================================================



# Main data set ----------------------------------------------------------------
DATA <- as.data.frame(DATA)

# Placebo data -----------------------------------------------------------------
DATA.plac <- as.data.frame(DATA.plac)



# Descriptive analysis - Panel view---------------------------------------------

panelview(
  rel_red ~ TreatNZL.97,
  data = filter(DATA, !(ccode %in% c("OECD", "IRL")) &
                  (year >= 1960 & year <= 2017)),
  index = c("ccode", "year"),
  xlab = "Year",
  ylab = "Country",
  axis.adjust = T,
  theme.bw = TRUE,
  pre.post = TRUE,
  #  type = "outcome",
  by.group = F,
  cex.main = 11,
  cex.main.sub = 11,
  legendOff = F,
  #  color = c("#bbbbbb", "#777777", "blue"),
)

# Descriptive analysis - Plot trends -------------------------------------------

ggplot(data = filter(DATA,!(ccode %in% c("IRL")) &
                       (year >= 1970 & year <= 2017))) +
  geom_line(aes(x = year, y = rel_red, color = ccode)) + theme_bw() + theme(legend.position = "bottom") + labs(colour = "Country") + theme_pubr() + ylab("UESKGEN")



#####################
#### Modelling ######
#####################

# DM-LFM MODELS ================================================================

# Default seeds used in modelling: 1234 ----------------------------------------

# I. EFFPARLEG models ==========================================================

set.seed(1234) # 0, 10, 50, 1234, 5678
# bcC_out.s0 <- bpCausal(
# bcC_out.s10 <- bpCausal(
# bcC_out.s50 <- bpCausal(
bcC_out.s1234 <- bpCausal(
  # bcC_out.s5678 <- bpCausal(
  data = filter(DATA, !(ccode %in% c("OECD", "IRL")) &
                  (year <= 2017 & year >= 1960)),
  index = c("ccode", "year"),
  Yname = "effpar_leg",
  Dname = "TreatNZL.96",
  #  Dname = "TreatNZL.91",
  Xname = c(
    "inflation_cpi",
    "elderly",
    "young",
    "vp_rigidity",
    "eneg",
    "enpres.2",
    "pres",
    "realgdpgr",
    "fed",
    "bic",
    "popul",
    "instcons"
  ),
  Zname = c(
    "inflation_cpi",
    "elderly",
    "young",
    "vp_rigidity",
    "eneg",
    "enpres.2",
    "pres",
    "realgdpgr",
    "fed",
    "bic",
    "popul",
    "instcons"
  ),
  Aname = c(
    "inflation_cpi",
    "elderly",
    "young",
    "vp_rigidity",
    "eneg",
    "enpres.2",
    "pres",
    "realgdpgr",
    "fed",
    "bic",
    "popul",
    "instcons"
  ),
  re = "both",
  ar1 = T,
  r = 3,
  # r = 3
  niter = 100000,
  xlasso = 1,
  ## whether to shrink constant coefs (1 = TRUE, 0 = FALSE)
  zlasso = 1,
  ## whether to shrink unit-level random coefs (1 = TRUE, 0 = FALSE)
  alasso = 1,
  ## whether to shrink time-level coefs (1 = TRUE, 0 = FALSE)
  flasso = 1,
  ## whether to shrink factor loadings (1 = TRUE, 0 = FALSE)
  a1 = 0.001,
  a2 = 0.001,
  ## parameters for hyper prior shrink on beta (diffuse hyper priors)
  b1 = 0.001,
  b2 = 0.001,
  ## parameters for hyper prior shrink on alpha_i
  c1 = 0.001,
  c2 = 0.001,
  ## parameters for hyper prior shrink on xi_t
  p1 = 0.001,
  p2 = 0.001 ## parameters for hyper prior shrink on factor terms
)


# Estimation results summary ===================================================

coef.bcC_out <-
  coefSummary(bcC_out.s1234)  ## summary estimated parameters

eff.bcC_out <- effSummary(
  bcC_out.s1234,
  ## summary treatment effects
  usr.id = NULL,
  ## treatment effect for individual treated units, if input NULL, calculate average TT
  cumu = FALSE,
  ## whether to calculate culmulative treatment effects
  rela.period = TRUE
) ## whether to use time relative to the occurence of treatment (1 is the first post-treatment period) or real period (like year 1998, 1999, ...)

#coef.bcC_out$est.beta
eff.bcC_out$est.eff[c(1, 2, 5, 8)]
mean(eff.bcC_out$est.eff[rownames(eff.bcC_out$est.eff) > 0, 5])


# Figures ======================================================================

df.bpC <-
  data.frame(eff.bcC_out$est.eff)

# BAND Plots -------------------------------------------------------------------

ggplot() + geom_line(
  data = filter(df.bpC),
  aes(time, observed),
  alpha = 1,
  linewidth = 1
) + geom_line(
  data = filter(df.bpC),
  aes(time, estimated_counterfactual),
  alpha = 1,
  linetype = "dashed",
  color = "blue",
  linewidth = 1
) + geom_line(
  data = filter(df.bpC),
  aes(time, counterfactual_ci_u),
  alpha = 0.6,
  linetype = "dashed",
  color = "red",
  linewidth = 0.3
) + geom_line(
  data = filter(df.bpC),
  aes(time, counterfactual_ci_l),
  alpha = 0.6,
  linetype = "dashed",
  color = "red",
  linewidth = 0.3
) +
  theme_linedraw() +
  geom_vline(
    xintercept = c(0, 0),
    #    xintercept = c(0, 9),
    #    xintercept = c(0, 5),
    linetype = "dashed",
    color = "black",
    linewidth = 0.5
  ) + theme(legend.position = "none") + xlab("Time Relative to Treatment") + ylab("Treated vs. Synthetic Unit") + theme_pubr() + annotate(
    "rect",
    xmin = 0,
    xmax = 0,
    #    xmax = 9,
    #    xmax = 5,
    ymin = -Inf,
    ymax = +Inf,
    alpha = .2,
    fill = "green"
  )


# Gap plots --------------------------------------------------------------------

ggplot() + geom_line(
  data = filter(df.bpC),
  aes(time, estimated_ATT),
  alpha = 1,
  linewidth = 1
) + geom_line(
  data = filter(df.bpC),
  aes(time, estimated_ATT_ci_u),
  alpha = 0.6,
  color = "red",
  linewidth = 0.5,
  linetype = "dashed"
) + geom_line(
  data = filter(df.bpC),
  aes(time, estimated_ATT_ci_l),
  alpha = 0.6,
  color = "red",
  linewidth = 0.5,
  linetype = "dashed"
) +
  theme_linedraw() +
  geom_vline(
    xintercept = c(0, 0),
    # 5 or 9
    linetype = "dashed",
    color = "black",
    linewidth = 0.5
  ) +
  geom_hline(
    yintercept = 0,
    linetype = "dashed",
    color = "black",
    linewidth = 0.5
  ) +
  theme(legend.position = "none") + xlab("Time Relative to Treatment") + ylab("Treatment Effect") + theme_pubr() + annotate(
    "rect",
    xmin = 0,
    xmax = 0,
    # 5 or 9
    ymin = -Inf,
    ymax = +Inf,
    alpha = .2,
    fill = "green"
  )




# II. VTURN models =============================================================

set.seed(1234) # 0, 10, 50, 1234, 5678
# bcC_out.s0 <- bpCausal(
# bcC_out.s10 <- bpCausal(
# bcC_out.s50 <- bpCausal(
bcC_out.s1234 <- bpCausal(
  # bcC_out.s5678 <- bpCausal(
  data = filter(DATA, !(ccode %in% c("OECD", "IRL")) &
                  (year <= 2017 & year >= 1960)),
  index = c("ccode", "year"),
  Yname = "vturn",
  Dname = "TreatNZL.96",
  #  Dname = "TreatNZL.91",
  Xname = c(
    "inflation_cpi",
    "elderly",
    "young",
    "vp_rigidity",
    "eneg",
    "enpres.2",
    "pres",
    "realgdpgr",
    "fed",
    "bic",
    "popul",
    "instcons",
    "comp_voting"
  ),
  Zname = c(
    "inflation_cpi",
    "elderly",
    "young",
    "vp_rigidity",
    "eneg",
    "enpres.2",
    "pres",
    "realgdpgr",
    "fed",
    "bic",
    "popul",
    "instcons",
    "comp_voting"
  ),
  Aname = c(
    "inflation_cpi",
    "elderly",
    "young",
    "vp_rigidity",
    "eneg",
    "enpres.2",
    "pres",
    "realgdpgr",
    "fed",
    "bic",
    "popul",
    "instcons",
    "comp_voting"
  ),
  re = "both",
  ar1 = T,
  r = 6,
  # r = 6
  niter = 100000,
  xlasso = 1,
  ## whether to shrink constant coefs (1 = TRUE, 0 = FALSE)
  zlasso = 1,
  ## whether to shrink unit-level random coefs (1 = TRUE, 0 = FALSE)
  alasso = 1,
  ## whether to shrink time-level coefs (1 = TRUE, 0 = FALSE)
  flasso = 1,
  ## whether to shrink factor loadings (1 = TRUE, 0 = FALSE)
  a1 = 0.001,
  a2 = 0.001,
  ## parameters for hyper prior shrink on beta (diffuse hyper priors)
  b1 = 0.001,
  b2 = 0.001,
  ## parameters for hyper prior shrink on alpha_i
  c1 = 0.001,
  c2 = 0.001,
  ## parameters for hyper prior shrink on xi_t
  p1 = 0.001,
  p2 = 0.001 ## parameters for hyper prior shrink on factor terms
)

# Estimation results summary ===================================================

coef.bcC_out <-
  coefSummary(bcC_out.s1234)  ## summary estimated parameters

eff.bcC_out <- effSummary(
  bcC_out.s1234,
  ## summary treatment effects
  usr.id = NULL,
  ## treatment effect for individual treated units, if input NULL, calculate average TT
  cumu = FALSE,
  ## whether to calculate culmulative treatment effects
  rela.period = TRUE
) ## whether to use time relative to the occurence of treatment (1 is the first post-treatment period) or real period (like year 1998, 1999, ...)

#coef.bcC_out$est.beta
eff.bcC_out$est.eff[c(1, 2, 5, 8)]
mean(eff.bcC_out$est.eff[rownames(eff.bcC_out$est.eff) > 0, 5])


# Figures ======================================================================

df.bpC <-
  data.frame(eff.bcC_out$est.eff)

# BAND Plots -------------------------------------------------------------------

ggplot() + geom_line(
  data = filter(df.bpC),
  aes(time, observed),
  alpha = 1,
  linewidth = 1
) + geom_line(
  data = filter(df.bpC),
  aes(time, estimated_counterfactual),
  alpha = 1,
  linetype = "dashed",
  color = "blue",
  linewidth = 1
) + geom_line(
  data = filter(df.bpC),
  aes(time, counterfactual_ci_u),
  alpha = 0.6,
  linetype = "dashed",
  color = "red",
  linewidth = 0.3
) + geom_line(
  data = filter(df.bpC),
  aes(time, counterfactual_ci_l),
  alpha = 0.6,
  linetype = "dashed",
  color = "red",
  linewidth = 0.3
) +
  theme_linedraw() +
  geom_vline(
    xintercept = c(0, 0),
    #    xintercept = c(0, 9),
    #    xintercept = c(0, 5),
    linetype = "dashed",
    color = "black",
    linewidth = 0.5
  ) + theme(legend.position = "none") + xlab("Time Relative to Treatment") + ylab("Treated vs. Synthetic Unit") + theme_pubr() + annotate(
    "rect",
    xmin = 0,
    xmax = 0,
    #    xmax = 9,
    #    xmax = 5,
    ymin = -Inf,
    ymax = +Inf,
    alpha = .2,
    fill = "green"
  )


# Gap plots --------------------------------------------------------------------

ggplot() + geom_line(
  data = filter(df.bpC),
  aes(time, estimated_ATT),
  alpha = 1,
  linewidth = 1
) + geom_line(
  data = filter(df.bpC),
  aes(time, estimated_ATT_ci_u),
  alpha = 0.6,
  color = "red",
  linewidth = 0.5,
  linetype = "dashed"
) + geom_line(
  data = filter(df.bpC),
  aes(time, estimated_ATT_ci_l),
  alpha = 0.6,
  color = "red",
  linewidth = 0.5,
  linetype = "dashed"
) +
  theme_linedraw() +
  geom_vline(
    xintercept = c(0, 0),
    # 5 or 9
    linetype = "dashed",
    color = "black",
    linewidth = 0.5
  ) +
  geom_hline(
    yintercept = 0,
    linetype = "dashed",
    color = "black",
    linewidth = 0.5
  ) +
  theme(legend.position = "none") + xlab("Time Relative to Treatment") + ylab("Treatment Effect") + theme_pubr() + annotate(
    "rect",
    xmin = 0,
    xmax = 0,
    # 5 or 9
    ymin = -Inf,
    ymax = +Inf,
    alpha = .2,
    fill = "green"
  )



# |||. GOVRIGHT models =========================================================

set.seed(1234) # 0, 10, 50, 1234, 5678
#bcC_out.s0 <- bpCausal(
#bcC_out.s10 <- bpCausal(
#bcC_out.s50 <- bpCausal(
bcC_out.s1234 <- bpCausal(
  #bcC_out.s5678 <- bpCausal(
  data = filter(DATA, !(ccode %in% c("OECD", "IRL")) &
                  (year <= 2017 & year >= 1960)),
  index = c("ccode", "year"),
  Yname = "logit_gr1",
  Dname = "TreatNZL.96",
  #  Dname = "TreatNZL.91",
  Xname = c(
    "inflation_cpi",
    "elderly",
    "young",
    "vp_rigidity",
    "eneg",
    "enpres.2",
    "pres",
    "realgdpgr",
    "fed",
    "bic",
    "popul",
    "instcons"
  ),
  Zname = c(
    "inflation_cpi",
    "elderly",
    "young",
    "vp_rigidity",
    "eneg",
    "enpres.2",
    "pres",
    "realgdpgr",
    "fed",
    "bic",
    "popul",
    "instcons"
  ),
  Aname = c(
    "inflation_cpi",
    "elderly",
    "young",
    "vp_rigidity",
    "eneg",
    "enpres.2",
    "pres",
    "realgdpgr",
    "fed",
    "bic",
    "popul",
    "instcons"
  ),
  re = "both",
  ar1 = T,
  r = 1,
  # r = 1
  niter = 100000,
  xlasso = 1,
  ## whether to shrink constant coefs (1 = TRUE, 0 = FALSE)
  zlasso = 1,
  ## whether to shrink unit-level random coefs (1 = TRUE, 0 = FALSE)
  alasso = 1,
  ## whether to shrink time-level coefs (1 = TRUE, 0 = FALSE)
  flasso = 1,
  ## whether to shrink factor loadings (1 = TRUE, 0 = FALSE)
  a1 = 0.001,
  a2 = 0.001,
  ## parameters for hyper prior shrink on beta (diffuse hyper priors)
  b1 = 0.001,
  b2 = 0.001,
  ## parameters for hyper prior shrink on alpha_i
  c1 = 0.001,
  c2 = 0.001,
  ## parameters for hyper prior shrink on xi_t
  p1 = 0.001,
  p2 = 0.001 ## parameters for hyper prior shrink on factor terms
)

# Estimation results summary ===================================================

coef.bcC_out <-
  coefSummary(bcC_out.s1234)  ## summary estimated parameters

eff.bcC_out <- effSummary(
  bcC_out.s1234,
  ## summary treatment effects
  usr.id = NULL,
  ## treatment effect for individual treated units, if input NULL, calculate average TT
  cumu = FALSE,
  ## whether to calculate culmulative treatment effects
  rela.period = TRUE
) ## whether to use time relative to the occurence of treatment (1 is the first post-treatment period) or real period (like year 1998, 1999, ...)

#coef.bcC_out$est.beta
eff.bcC_out$est.eff[c(1, 2, 5, 8)]
mean(eff.bcC_out$est.eff[rownames(eff.bcC_out$est.eff) > 0, 5])


# Figures ======================================================================

df.bpC <-
  data.frame(eff.bcC_out$est.eff)

# BAND Plots -------------------------------------------------------------------

ggplot() + geom_line(
  data = filter(df.bpC),
  aes(time, observed),
  alpha = 1,
  linewidth = 1
) + geom_line(
  data = filter(df.bpC),
  aes(time, estimated_counterfactual),
  alpha = 1,
  linetype = "dashed",
  color = "blue",
  linewidth = 1
) + geom_line(
  data = filter(df.bpC),
  aes(time, counterfactual_ci_u),
  alpha = 0.6,
  linetype = "dashed",
  color = "red",
  linewidth = 0.3
) + geom_line(
  data = filter(df.bpC),
  aes(time, counterfactual_ci_l),
  alpha = 0.6,
  linetype = "dashed",
  color = "red",
  linewidth = 0.3
) +
  theme_linedraw() +
  geom_vline(
    xintercept = c(0, 0),
    #    xintercept = c(0, 9),
    #    xintercept = c(0, 5),
    linetype = "dashed",
    color = "black",
    linewidth = 0.5
  ) + theme(legend.position = "none") + xlab("Time Relative to Treatment") + ylab("Treated vs. Synthetic Unit") + theme_pubr() + annotate(
    "rect",
    xmin = 0,
    xmax = 0,
    #    xmax = 9,
    #    xmax = 5,
    ymin = -Inf,
    ymax = +Inf,
    alpha = .2,
    fill = "green"
  )


# Gap plots --------------------------------------------------------------------

ggplot() + geom_line(
  data = filter(df.bpC),
  aes(time, estimated_ATT),
  alpha = 1,
  linewidth = 1
) + geom_line(
  data = filter(df.bpC),
  aes(time, estimated_ATT_ci_u),
  alpha = 0.6,
  color = "red",
  linewidth = 0.5,
  linetype = "dashed"
) + geom_line(
  data = filter(df.bpC),
  aes(time, estimated_ATT_ci_l),
  alpha = 0.6,
  color = "red",
  linewidth = 0.5,
  linetype = "dashed"
) +
  theme_linedraw() +
  geom_vline(
    xintercept = c(0, 0),
    # 5 or 9
    linetype = "dashed",
    color = "black",
    linewidth = 0.5
  ) +
  geom_hline(
    yintercept = 0,
    linetype = "dashed",
    color = "black",
    linewidth = 0.5
  ) +
  theme(legend.position = "none") + xlab("Time Relative to Treatment") + ylab("Treatment Effect") + theme_pubr() + annotate(
    "rect",
    xmin = 0,
    xmax = 0,
    # 5 or 9
    ymin = -Inf,
    ymax = +Inf,
    alpha = .2,
    fill = "green"
  )


# IV. SCONSERV models ==========================================================

set.seed(1234) # 0, 10, 50, 1234, 5678
#bcC_out.s0 <- bpCausal(
#bcC_out.s10 <- bpCausal(
#bcC_out.s50 <- bpCausal(
bcC_out.s1234 <- bpCausal(
  #bcC_out.s5678 <- bpCausal(
  data = filter(DATA, !(ccode %in% c("OECD", "IRL")) &
                  (year <= 2017 & year >= 1960)),
  index = c("ccode", "year"),
  # Yname = "sconserv",
  Yname = "logit.sconserv",
  Dname = "TreatNZL.96",
  #  Dname = "TreatNZL.91",
  Xname = c(
    "inflation_cpi",
    "elderly",
    "young",
    "vp_rigidity",
    "eneg",
    "enpres.2",
    "pres",
    "realgdpgr",
    "fed",
    "bic",
    "popul",
    "instcons"
  ),
  Zname = c(
    "inflation_cpi",
    "elderly",
    "young",
    "vp_rigidity",
    "eneg",
    "enpres.2",
    "pres",
    "realgdpgr",
    "fed",
    "bic",
    "popul",
    "instcons"
  ),
  Aname = c(
    "inflation_cpi",
    "elderly",
    "young",
    "vp_rigidity",
    "eneg",
    "enpres.2",
    "pres",
    "realgdpgr",
    "fed",
    "bic",
    "popul",
    "instcons"
  ),
  re = "both",
  ar1 = T,
  r = 3,
  # r = 3 (logit.sconserv)
  # r = 1 (sconserv)
  niter = 100000,
  xlasso = 1,
  ## whether to shrink constant coefs (1 = TRUE, 0 = FALSE)
  zlasso = 1,
  ## whether to shrink unit-level random coefs (1 = TRUE, 0 = FALSE)
  alasso = 1,
  ## whether to shrink time-level coefs (1 = TRUE, 0 = FALSE)
  flasso = 1,
  ## whether to shrink factor loadings (1 = TRUE, 0 = FALSE)
  a1 = 0.001,
  a2 = 0.001,
  ## parameters for hyper prior shrink on beta (diffuse hyper priors)
  b1 = 0.001,
  b2 = 0.001,
  ## parameters for hyper prior shrink on alpha_i
  c1 = 0.001,
  c2 = 0.001,
  ## parameters for hyper prior shrink on xi_t
  p1 = 0.001,
  p2 = 0.001 ## parameters for hyper prior shrink on factor terms
)

# Estimation results summary ===================================================

coef.bcC_out <-
  coefSummary(bcC_out.s1234)  ## summary estimated parameters

eff.bcC_out <- effSummary(
  bcC_out.s1234,
  ## summary treatment effects
  usr.id = NULL,
  ## treatment effect for individual treated units, if input NULL, calculate average TT
  cumu = FALSE,
  ## whether to calculate culmulative treatment effects
  rela.period = TRUE
) ## whether to use time relative to the occurence of treatment (1 is the first post-treatment period) or real period (like year 1998, 1999, ...)

#coef.bcC_out$est.beta
eff.bcC_out$est.eff[c(1, 2, 5, 8)]
mean(eff.bcC_out$est.eff[rownames(eff.bcC_out$est.eff) > 0, 5])


# Figures ======================================================================

df.bpC <-
  data.frame(eff.bcC_out$est.eff)

# BAND Plots -------------------------------------------------------------------

ggplot() + geom_line(
  data = filter(df.bpC),
  aes(time, observed),
  alpha = 1,
  linewidth = 1
) + geom_line(
  data = filter(df.bpC),
  aes(time, estimated_counterfactual),
  alpha = 1,
  linetype = "dashed",
  color = "blue",
  linewidth = 1
) + geom_line(
  data = filter(df.bpC),
  aes(time, counterfactual_ci_u),
  alpha = 0.6,
  linetype = "dashed",
  color = "red",
  linewidth = 0.3
) + geom_line(
  data = filter(df.bpC),
  aes(time, counterfactual_ci_l),
  alpha = 0.6,
  linetype = "dashed",
  color = "red",
  linewidth = 0.3
) +
  theme_linedraw() +
  geom_vline(
    xintercept = c(0, 0),
    #    xintercept = c(0, 9),
    #    xintercept = c(0, 5),
    linetype = "dashed",
    color = "black",
    linewidth = 0.5
  ) + theme(legend.position = "none") + xlab("Time Relative to Treatment") + ylab("Treated vs. Synthetic Unit") + theme_pubr() + annotate(
    "rect",
    xmin = 0,
    xmax = 0,
    #    xmax = 9,
    #    xmax = 5,
    ymin = -Inf,
    ymax = +Inf,
    alpha = .2,
    fill = "green"
  )


# Gap plots --------------------------------------------------------------------

ggplot() + geom_line(
  data = filter(df.bpC),
  aes(time, estimated_ATT),
  alpha = 1,
  linewidth = 1
) + geom_line(
  data = filter(df.bpC),
  aes(time, estimated_ATT_ci_u),
  alpha = 0.6,
  color = "red",
  linewidth = 0.5,
  linetype = "dashed"
) + geom_line(
  data = filter(df.bpC),
  aes(time, estimated_ATT_ci_l),
  alpha = 0.6,
  color = "red",
  linewidth = 0.5,
  linetype = "dashed"
) +
  theme_linedraw() +
  geom_vline(
    xintercept = c(0, 0),
    # 5 or 9
    linetype = "dashed",
    color = "black",
    linewidth = 0.5
  ) +
  geom_hline(
    yintercept = 0,
    linetype = "dashed",
    color = "black",
    linewidth = 0.5
  ) +
  theme(legend.position = "none") + xlab("Time Relative to Treatment") + ylab("Treatment Effect") + theme_pubr() + annotate(
    "rect",
    xmin = 0,
    xmax = 0,
    # 5 or 9
    ymin = -Inf,
    ymax = +Inf,
    alpha = .2,
    fill = "green"
  )


# V. RELRED models (including 'in-space' placebos for RELRED and UESKGEN)

set.seed(1234) # 0 (1), 10, 50, 1234, 5678
#bcC_out.s1 <- bpCausal(
#bcC_out.s10 <- bpCausal(
#bcC_out.s50 <- bpCausal(
bcC_out.s1234 <- bpCausal(
  #bcC_out.s5678 <- bpCausal(
  data = filter(DATA, !(ccode %in% c("OECD", "IRL")) &
                  (year <= 2017 & year >= 1960)),
  #  data = filter(DATA.plac,!(ccode %in% c("OECD", "IRL")) & (year <= 2017 & year >= 1960)),
  index = c("ccode", "year"),
  Yname = "rel_red",
  #   Dname = "TreatNZL.92",
  Dname = "TreatNZL.97",
  #   Dname = "TreatGBR.97",
  #   Dname = "TreatAUS.97",
  #   Dname = "TreatFRA.97",
  #   Dname = "TreatCAN.97",
  #   Dname = "TreatUSA.97",
  #   Dname = "TreatJPN.97",
  Xname = c(
    "inflation_cpi",
    "elderly",
    "rnetu",
    "young",
    "vp_rigidity",
    "eneg",
    "pres",
    "realgdpgr",
    "fed",
    "bic",
    "popul",
    "topone_inc"
  ),
  Zname = c(
    "inflation_cpi",
    "elderly",
    "rnetu",
    "young",
    "vp_rigidity",
    "eneg",
    "pres",
    "realgdpgr",
    "fed",
    "bic",
    "popul",
    "topone_inc"
  ),
  Aname = c(
    "inflation_cpi",
    "elderly",
    "rnetu",
    "young",
    "vp_rigidity",
    "eneg",
    "pres",
    "realgdpgr",
    "fed",
    "bic",
    "popul",
    "topone_inc"
  ),
  re = "both",
  ar1 = T,
  r = 2,
  # r = 2, (RELRED)
  niter = 100000,
  xlasso = 1,
  ## whether to shrink constant coefs (1 = TRUE, 0 = FALSE)
  zlasso = 1,
  ## whether to shrink unit-level random coefs (1 = TRUE, 0 = FALSE)
  alasso = 1,
  ## whether to shrink time-level coefs (1 = TRUE, 0 = FALSE)
  flasso = 1,
  ## whether to shrink factor loadings (1 = TRUE, 0 = FALSE)
  a1 = 0.001,
  a2 = 0.001,
  ## parameters for hyper prior shrink on beta (diffuse hyper priors)
  b1 = 0.001,
  b2 = 0.001,
  ## parameters for hyper prior shrink on alpha_i
  c1 = 0.001,
  c2 = 0.001,
  ## parameters for hyper prior shrink on xi_t
  p1 = 0.001,
  p2 = 0.001 ## parameters for hyper prior shrink on factor terms
)

# Estimation results summary ===================================================

coef.bcC_out <-
  coefSummary(bcC_out.s1234)  ## summary estimated parameters

eff.bcC_out <- effSummary(
  bcC_out.s1234,
  ## summary treatment effects
  usr.id = NULL,
  ## treatment effect for individual treated units, if input NULL, calculate average TT
  cumu = FALSE,
  ## whether to calculate culmulative treatment effects
  rela.period = TRUE
) ## whether to use time relative to the occurence of treatment (1 is the first post-treatment period) or real period (like year 1998, 1999, ...)

#coef.bcC_out$est.beta
eff.bcC_out$est.eff[c(1, 2, 5, 8)]
mean(eff.bcC_out$est.eff[rownames(eff.bcC_out$est.eff) > 0, 5])


# Figures ======================================================================

df.bpC <- data.frame(eff.bcC_out$est.eff)

# BAND Plots -------------------------------------------------------------------

ggplot() + geom_line(
  data = filter(df.bpC),
  aes(time, observed),
  alpha = 1,
  linewidth = 1
) + geom_line(
  data = filter(df.bpC),
  aes(time, estimated_counterfactual),
  alpha = 1,
  linetype = "dashed",
  color = "blue",
  linewidth = 1
) + geom_line(
  data = filter(df.bpC),
  aes(time, counterfactual_ci_u),
  alpha = 0.6,
  linetype = "dashed",
  color = "red",
  linewidth = 0.3
) + geom_line(
  data = filter(df.bpC),
  aes(time, counterfactual_ci_l),
  alpha = 0.6,
  linetype = "dashed",
  color = "red",
  linewidth = 0.3
) +
  theme_linedraw() +
  geom_vline(
    xintercept = c(0, 0),
    #    xintercept = c(0, 9),
    #    xintercept = c(0, 5),
    linetype = "dashed",
    color = "black",
    linewidth = 0.5
  ) + theme(legend.position = "none") + xlab("Time Relative to Treatment") + ylab("Treated vs. Synthetic Unit") + theme_pubr() + annotate(
    "rect",
    xmin = 0,
    xmax = 0,
    #    xmax = 9,
    #    xmax = 5,
    ymin = -Inf,
    ymax = +Inf,
    alpha = .2,
    fill = "green"
  )


# Gap plots --------------------------------------------------------------------

ggplot() + geom_line(
  data = filter(df.bpC),
  aes(time, estimated_ATT),
  alpha = 1,
  linewidth = 1
) + geom_line(
  data = filter(df.bpC),
  aes(time, estimated_ATT_ci_u),
  alpha = 0.6,
  color = "red",
  linewidth = 0.5,
  linetype = "dashed"
) + geom_line(
  data = filter(df.bpC),
  aes(time, estimated_ATT_ci_l),
  alpha = 0.6,
  color = "red",
  linewidth = 0.5,
  linetype = "dashed"
) +
  theme_linedraw() +
  geom_vline(
    xintercept = c(0, 0),
    # 5 or 9
    linetype = "dashed",
    color = "black",
    linewidth = 0.5
  ) +
  geom_hline(
    yintercept = 0,
    linetype = "dashed",
    color = "black",
    linewidth = 0.5
  ) +
  theme(legend.position = "none") + xlab("Time Relative to Treatment") + ylab("Treatment Effect") + theme_pubr() + annotate(
    "rect",
    xmin = 0,
    xmax = 0,
    # 5 or 9
    ymin = -Inf,
    ymax = +Inf,
    alpha = .2,
    fill = "green"
  )


# VI. UESKGEN models (including 'in-space' placebos for RELRED and UESKGEN)

set.seed(1234) # 0 (1), 10, 50, 1234, 5678
#bcC_out.s1 <- bpCausal(
#bcC_out.s10 <- bpCausal(
#bcC_out.s50 <- bpCausal(
bcC_out.s1234 <- bpCausal(
  #bcC_out.s5678 <- bpCausal(
  data = filter(DATA, !(ccode %in% c("OECD", "IRL")) &
                  (year <= 2017 & year >= 1960)),
  #  data = filter(DATA.plac,!(ccode %in% c("OECD", "IRL")) & (year <= 2017 & year >= 1960)),
  index = c("ccode", "year"),
  Yname = "UESKGEN",
  #   Dname = "TreatNZL.92",
  Dname = "TreatNZL.97",
  #   Dname = "TreatGBR.97",
  #   Dname = "TreatAUS.97",
  #   Dname = "TreatFRA.97",
  #   Dname = "TreatCAN.97",
  #   Dname = "TreatUSA.97",
  #   Dname = "TreatJPN.97",
  Xname = c(
    "inflation_cpi",
    "elderly",
    "rnetu",
    "young",
    "vp_rigidity",
    "eneg",
    "pres",
    "realgdpgr",
    "fed",
    "bic",
    "popul",
    "topone_inc"
  ),
  Zname = c(
    "inflation_cpi",
    "elderly",
    "rnetu",
    "young",
    "vp_rigidity",
    "eneg",
    "pres",
    "realgdpgr",
    "fed",
    "bic",
    "popul",
    "topone_inc"
  ),
  Aname = c(
    "inflation_cpi",
    "elderly",
    "rnetu",
    "young",
    "vp_rigidity",
    "eneg",
    "pres",
    "realgdpgr",
    "fed",
    "bic",
    "popul",
    "topone_inc"
  ),
  re = "both",
  ar1 = T,
  r = 3,
  # r = 3 (UESKGEN)
  niter = 100000,
  xlasso = 1,
  ## whether to shrink constant coefs (1 = TRUE, 0 = FALSE)
  zlasso = 1,
  ## whether to shrink unit-level random coefs (1 = TRUE, 0 = FALSE)
  alasso = 1,
  ## whether to shrink time-level coefs (1 = TRUE, 0 = FALSE)
  flasso = 1,
  ## whether to shrink factor loadings (1 = TRUE, 0 = FALSE)
  a1 = 0.001,
  a2 = 0.001,
  ## parameters for hyper prior shrink on beta (diffuse hyper priors)
  b1 = 0.001,
  b2 = 0.001,
  ## parameters for hyper prior shrink on alpha_i
  c1 = 0.001,
  c2 = 0.001,
  ## parameters for hyper prior shrink on xi_t
  p1 = 0.001,
  p2 = 0.001 ## parameters for hyper prior shrink on factor terms
)

# Estimation results summary ===================================================

coef.bcC_out <-
  coefSummary(bcC_out.s1234)  ## summary estimated parameters

eff.bcC_out <- effSummary(
  bcC_out.s1234,
  ## summary treatment effects
  usr.id = NULL,
  ## treatment effect for individual treated units, if input NULL, calculate average TT
  cumu = FALSE,
  ## whether to calculate culmulative treatment effects
  rela.period = TRUE
) ## whether to use time relative to the occurence of treatment (1 is the first post-treatment period) or real period (like year 1998, 1999, ...)

#coef.bcC_out$est.beta
eff.bcC_out$est.eff[c(1, 2, 5, 8)]
mean(eff.bcC_out$est.eff[rownames(eff.bcC_out$est.eff) > 0, 5])


# Figures ======================================================================

df.bpC <- data.frame(eff.bcC_out$est.eff)

# BAND Plots -------------------------------------------------------------------

ggplot() + geom_line(
  data = filter(df.bpC),
  aes(time, observed),
  alpha = 1,
  linewidth = 1
) + geom_line(
  data = filter(df.bpC),
  aes(time, estimated_counterfactual),
  alpha = 1,
  linetype = "dashed",
  color = "blue",
  linewidth = 1
) + geom_line(
  data = filter(df.bpC),
  aes(time, counterfactual_ci_u),
  alpha = 0.6,
  linetype = "dashed",
  color = "red",
  linewidth = 0.3
) + geom_line(
  data = filter(df.bpC),
  aes(time, counterfactual_ci_l),
  alpha = 0.6,
  linetype = "dashed",
  color = "red",
  linewidth = 0.3
) +
  theme_linedraw() +
  geom_vline(
    xintercept = c(0, 0),
    #    xintercept = c(0, 9),
    #    xintercept = c(0, 5),
    linetype = "dashed",
    color = "black",
    linewidth = 0.5
  ) + theme(legend.position = "none") + xlab("Time Relative to Treatment") + ylab("Treated vs. Synthetic Unit") + theme_pubr() + annotate(
    "rect",
    xmin = 0,
    xmax = 0,
    #    xmax = 9,
    #    xmax = 5,
    ymin = -Inf,
    ymax = +Inf,
    alpha = .2,
    fill = "green"
  )


# Gap plots --------------------------------------------------------------------

ggplot() + geom_line(
  data = filter(df.bpC),
  aes(time, estimated_ATT),
  alpha = 1,
  linewidth = 1
) + geom_line(
  data = filter(df.bpC),
  aes(time, estimated_ATT_ci_u),
  alpha = 0.6,
  color = "red",
  linewidth = 0.5,
  linetype = "dashed"
) + geom_line(
  data = filter(df.bpC),
  aes(time, estimated_ATT_ci_l),
  alpha = 0.6,
  color = "red",
  linewidth = 0.5,
  linetype = "dashed"
) +
  theme_linedraw() +
  geom_vline(
    xintercept = c(0, 0),
    # 5 or 9
    linetype = "dashed",
    color = "black",
    linewidth = 0.5
  ) +
  geom_hline(
    yintercept = 0,
    linetype = "dashed",
    color = "black",
    linewidth = 0.5
  ) +
  theme(legend.position = "none") + xlab("Time Relative to Treatment") + ylab("Treatment Effect") + theme_pubr() + annotate(
    "rect",
    xmin = 0,
    xmax = 0,
    # 5 or 9
    ymin = -Inf,
    ymax = +Inf,
    alpha = .2,
    fill = "green"
  )


# VII. GOVEXP models ===========================================================

set.seed(1234) # 0 (1), 10, 50, 1234, 5678
#bcC_out.s1 <- bpCausal(
#bcC_out.s10 <- bpCausal(
#bcC_out.s50 <- bpCausal(
bcC_out.s1234 <- bpCausal(
  #bcC_out.s5678 <- bpCausal(
  data = filter(DATA, !(ccode %in% c("OECD", "IRL")) &
                  (year <= 2017 & year >= 1960)),
  #  data = filter(DATA.plac,!(ccode %in% c("OECD", "IRL")) & (year <= 2017 & year >= 1960)),
  index = c("ccode", "year"),
  Yname = "govexp",
  #   Dname = "TreatNZL.92",
  Dname = "TreatNZL.97",
  #   Dname = "TreatGBR.97",
  #   Dname = "TreatAUS.97",
  #   Dname = "TreatFRA.97",
  #   Dname = "TreatCAN.97",
  #   Dname = "TreatUSA.97",
  #   Dname = "TreatJPN.97",
  Xname = c(
    "inflation_cpi",
    "elderly",
    "rnetu",
    "young",
    "vp_rigidity",
    "eneg",
    "pres",
    "realgdpgr",
    "fed",
    "bic",
    "popul",
    "topone_inc"
  ),
  Zname = c(
    "inflation_cpi",
    "elderly",
    "rnetu",
    "young",
    "vp_rigidity",
    "eneg",
    "pres",
    "realgdpgr",
    "fed",
    "bic",
    "popul",
    "topone_inc"
  ),
  Aname = c(
    "inflation_cpi",
    "elderly",
    "rnetu",
    "young",
    "vp_rigidity",
    "eneg",
    "pres",
    "realgdpgr",
    "fed",
    "bic",
    "popul",
    "topone_inc"
  ),
  re = "both",
  ar1 = T,
  r = 4,
  # r = 4 (GOVEXP)
  niter = 100000,
  xlasso = 1,
  ## whether to shrink constant coefs (1 = TRUE, 0 = FALSE)
  zlasso = 1,
  ## whether to shrink unit-level random coefs (1 = TRUE, 0 = FALSE)
  alasso = 1,
  ## whether to shrink time-level coefs (1 = TRUE, 0 = FALSE)
  flasso = 1,
  ## whether to shrink factor loadings (1 = TRUE, 0 = FALSE)
  a1 = 0.001,
  a2 = 0.001,
  ## parameters for hyper prior shrink on beta (diffuse hyper priors)
  b1 = 0.001,
  b2 = 0.001,
  ## parameters for hyper prior shrink on alpha_i
  c1 = 0.001,
  c2 = 0.001,
  ## parameters for hyper prior shrink on xi_t
  p1 = 0.001,
  p2 = 0.001 ## parameters for hyper prior shrink on factor terms
)

# Estimation results summary ===================================================

coef.bcC_out <-
  coefSummary(bcC_out.s1234)  ## summary estimated parameters

eff.bcC_out <- effSummary(
  bcC_out.s1234,
  ## summary treatment effects
  usr.id = NULL,
  ## treatment effect for individual treated units, if input NULL, calculate average TT
  cumu = FALSE,
  ## whether to calculate culmulative treatment effects
  rela.period = TRUE
) ## whether to use time relative to the occurence of treatment (1 is the first post-treatment period) or real period (like year 1998, 1999, ...)

#coef.bcC_out$est.beta
eff.bcC_out$est.eff[c(1, 2, 5, 8)]
mean(eff.bcC_out$est.eff[rownames(eff.bcC_out$est.eff) > 0, 5])


# Figures ======================================================================

df.bpC <- data.frame(eff.bcC_out$est.eff)

# BAND Plots -------------------------------------------------------------------

ggplot() + geom_line(
  data = filter(df.bpC),
  aes(time, observed),
  alpha = 1,
  linewidth = 1
) + geom_line(
  data = filter(df.bpC),
  aes(time, estimated_counterfactual),
  alpha = 1,
  linetype = "dashed",
  color = "blue",
  linewidth = 1
) + geom_line(
  data = filter(df.bpC),
  aes(time, counterfactual_ci_u),
  alpha = 0.6,
  linetype = "dashed",
  color = "red",
  linewidth = 0.3
) + geom_line(
  data = filter(df.bpC),
  aes(time, counterfactual_ci_l),
  alpha = 0.6,
  linetype = "dashed",
  color = "red",
  linewidth = 0.3
) +
  theme_linedraw() +
  geom_vline(
    xintercept = c(0, 0),
    #    xintercept = c(0, 9),
    #    xintercept = c(0, 5),
    linetype = "dashed",
    color = "black",
    linewidth = 0.5
  ) + theme(legend.position = "none") + xlab("Time Relative to Treatment") + ylab("Treated vs. Synthetic Unit") + theme_pubr() + annotate(
    "rect",
    xmin = 0,
    xmax = 0,
    #    xmax = 9,
    #    xmax = 5,
    ymin = -Inf,
    ymax = +Inf,
    alpha = .2,
    fill = "green"
  )


# Gap plots --------------------------------------------------------------------

ggplot() + geom_line(
  data = filter(df.bpC),
  aes(time, estimated_ATT),
  alpha = 1,
  linewidth = 1
) + geom_line(
  data = filter(df.bpC),
  aes(time, estimated_ATT_ci_u),
  alpha = 0.6,
  color = "red",
  linewidth = 0.5,
  linetype = "dashed"
) + geom_line(
  data = filter(df.bpC),
  aes(time, estimated_ATT_ci_l),
  alpha = 0.6,
  color = "red",
  linewidth = 0.5,
  linetype = "dashed"
) +
  theme_linedraw() +
  geom_vline(
    xintercept = c(0, 0),
    # 5 or 9
    linetype = "dashed",
    color = "black",
    linewidth = 0.5
  ) +
  geom_hline(
    yintercept = 0,
    linetype = "dashed",
    color = "black",
    linewidth = 0.5
  ) +
  theme(legend.position = "none") + xlab("Time Relative to Treatment") + ylab("Treatment Effect") + theme_pubr() + annotate(
    "rect",
    xmin = 0,
    xmax = 0,
    # 5 or 9
    ymin = -Inf,
    ymax = +Inf,
    alpha = .2,
    fill = "green"
  )




###################
### Diagnostics ###
###################


# Latent factors selection =====================================================

# Example: if three factors are included ---------------------------------------
wg.t <-
  data.frame(F1 = bcC_out.s1234$wg[1, ],
             F2 = bcC_out.s1234$wg[2, ],
             F3 = bcC_out.s1234$wg[3, ])


# Plot factors -----------------------------------------------------------------

wg.t.longer <- wg.t %>%
  pivot_longer(
    cols = starts_with("F"),
    names_to = "Factor",
    values_to = "loadings",
    values_drop_na = TRUE
  )

ggpubr::ggdensity(
  wg.t.longer,
  x = "loadings",
  facet.by = "Factor",
  fill = "yellow",
  add = "mean",
  rug = TRUE
) + theme_pubr()



# Diagnostics on MCMC ==========================================================

# MCMC convergence diagnostic (R functions retrieved from Pang et al. (2021) https://doi.org/10.7910/DVN/B6SWA1)


# Functions --------------------------------------------------------------------

## plot trace for att !!!
att.trace <- function(x) {
  pos <- which(x$rela.time.tr >= 0) ## post treatment period
  yo <- x$yo_t[pos] ## observed
  yct <- x$yct[pos,]
  niter <- dim(yct)[2]
  att <- sapply(1:niter, function(i) {
    return(mean(yo - yct[, i]))
  })
  return(att)
}

## plot trace for yct
yct.trace <- function(x) {
  pos <- which(x$rela.time.tr >= 0) ## post treatment period
  yct <- x$yct[pos,]
  niter <- dim(yct)[2]
  yct <- sapply(1:niter, function(i) {
    return(mean(yct[, i]))
  })
  return(yct)
}

## plot trace for yct (pre- and post-treatment period)
yct.trace.2 <- function(x) {
  yct <- x$yct
  niter <- dim(yct)[2]
  yct <- sapply(1:niter, function(i) {
    return(mean(yct[, i]))
  })
  return(yct)
}

## itt trace
itt.trace <- function(x, unit, time) {
  pos <- which(x$rela.time.tr == 5)[unit] ## post treatment period
  yo <- x$yo_t[pos] ## observed
  yct <- x$yct[pos,]
  niter <- dim(yct)[2]
  tt <- yo - yct
  return(tt)
}


## plot MCMC trace

plot.trace <- function(x, xlab, ylab) {
  data <-
    cbind.data.frame(x, 1:length(x), as.factor(rep(1, length(x))))
  names(data) <- c("trace", "time", "chain")
  p <- ggplot(data) + xlab(xlab) + ylab(ylab)
  p <- p + geom_line(aes(time, trace,
                         colour = chain,
                         group = chain), colour = "blue")
  p <- p + theme_pubr()
  return(p)
}

# ------------------------------------------------------------------------------


# Trace plots for ATT (post-treatment) =========================================

## att trace (effparleg)
att_effparleg.1 <- att.trace(bcC_out.s0)
att_effparleg.2 <- att.trace(bcC_out.s10)
att_effparleg.3 <- att.trace(bcC_out.s50)
att_effparleg.4 <- att.trace(bcC_out.s1234)
att_effparleg.5 <- att.trace(bcC_out.s5678)

## att trace (vturn)
att_vturn.1 <- att.trace(bcC_out.s0)
att_vturn.2 <- att.trace(bcC_out.s10)
att_vturn.3 <- att.trace(bcC_out.s50)
att_vturn.4 <- att.trace(bcC_out.s1234)
att_vturn.5 <- att.trace(bcC_out.s5678)

## att trace (logit.govright)
att_govright.1 <- att.trace(bcC_out.2.s0)
att_govright.2 <- att.trace(bcC_out.2.s10)
att_govright.3 <- att.trace(bcC_out.2.s50)
att_govright.4 <- att.trace(bcC_out.2.s1234)
att_govright.5 <- att.trace(bcC_out.2.s5678)

## att trace (logit.sconserv)
att_sconserv.1 <- att.trace(bcC_out.2.s0)
att_sconserv.2 <- att.trace(bcC_out.2.s10)
att_sconserv.3 <- att.trace(bcC_out.2.s50)
att_sconserv.4 <- att.trace(bcC_out.2.s1234)
att_sconserv.5 <- att.trace(bcC_out.2.s5678)

## att trace (relred, ...)

att_relred.1 <- att.trace(bcC_out.3.0)
att_relred.2 <- att.trace(bcC_out.3.10)
att_relred.3 <- att.trace(bcC_out.3.50)
att_relred.4 <- att.trace(bcC_out.3.1234)
att_relred.5 <- att.trace(bcC_out.3.5678)

att_ueskgen.1 <- att.trace(bcC_out.3.0)
att_ueskgen.2 <- att.trace(bcC_out.3.10)
att_ueskgen.3 <- att.trace(bcC_out.3.50)
att_ueskgen.4 <- att.trace(bcC_out.3.1234)
att_ueskgen.5 <- att.trace(bcC_out.3.5678)

att_govexp.1 <- att.trace(bcC_out.3.1)
att_govexp.2 <- att.trace(bcC_out.3.10)
att_govexp.3 <- att.trace(bcC_out.3.50)
att_govexp.4 <- att.trace(bcC_out.3.1234)
att_govexp.5 <- att.trace(bcC_out.3.5678)

# ==============================================================================


# Plots ----------------------------------------------------------

# Example: Trace plot for effparleg.1 (model for effparleg, seed = 0)
trace.att <- plot.trace(att_effparleg.1, "Iterations", "ATT")
trace.att


# yct <- yct.trace.2(bcC_out.s1234)
# itt <- itt.trace(bcC_out.s1234, 1, 5)  # unit 1, period 5
# trace.yct <- plot.trace(yct, "Iterations", "YCT", c(1.8, 2.2))
# trace.yct
# trace.itt <- plot.trace(itt, "Iterations", "TE(1, 25)", c(0, 11))
# trace.itt

# ------------------------------------------------------------------------------


# Gelman-Rubin convergence tests ===============================================

# Gelman-Rubin

# The `potential scale reduction factor' (PSRF) is an estimated factor by which the scale of the current distribution for the target distribution might be reduced if the simulations were continued for an infinite number of iterations.


# effparleg

mcmc.att_effparleg.1 <- mcmc(att_effparleg.1)
mcmc.att_effparleg.2 <- mcmc(att_effparleg.2)
mcmc.att_effparleg.3 <- mcmc(att_effparleg.3)
mcmc.att_effparleg.4 <- mcmc(att_effparleg.4)
mcmc.att_effparleg.5 <- mcmc(att_effparleg.5)

mcmc_list <-
  mcmc.list(
    mcmc.att_effparleg.1,
    mcmc.att_effparleg.2,
    mcmc.att_effparleg.3,
    mcmc.att_effparleg.4,
    mcmc.att_effparleg.5
  )

is.mcmc.list(mcmc_list)
gelman.diag(mcmc_list)
gelman.plot(mcmc_list)

# vturn

mcmc.att_effparleg.1 <- mcmc(att_effparleg.1)
mcmc.att_effparleg.2 <- mcmc(att_effparleg.2)
mcmc.att_effparleg.3 <- mcmc(att_effparleg.3)
mcmc.att_effparleg.4 <- mcmc(att_effparleg.4)
mcmc.att_effparleg.5 <- mcmc(att_effparleg.5)

mcmc_list <-
  mcmc.list(
    mcmc.att_effparleg.1,
    mcmc.att_effparleg.2,
    mcmc.att_effparleg.3,
    mcmc.att_effparleg.4,
    mcmc.att_effparleg.5
  )

is.mcmc.list(mcmc_list)
gelman.diag(mcmc_list)
gelman.plot(mcmc_list)

# logit.govright

mcmc.att_govright.1 <- mcmc(att_govright.1)
mcmc.att_govright.2 <- mcmc(att_govright.2)
mcmc.att_govright.3 <- mcmc(att_govright.3)
mcmc.att_govright.4 <- mcmc(att_govright.4)
mcmc.att_govright.5 <- mcmc(att_govright.5)

mcmc_list <-
  mcmc.list(
    mcmc.att_govright.1,
    mcmc.att_govright.2,
    mcmc.att_govright.3,
    mcmc.att_govright.4,
    mcmc.att_govright.5
  )

is.mcmc.list(mcmc_list)
gelman.diag(mcmc_list)
gelman.plot(mcmc_list)

# logit.sconserv

mcmc.att_sconserv.1 <- mcmc(att_sconserv.1)
mcmc.att_sconserv.2 <- mcmc(att_sconserv.2)
mcmc.att_sconserv.3 <- mcmc(att_sconserv.3)
mcmc.att_sconserv.4 <- mcmc(att_sconserv.4)
mcmc.att_sconserv.5 <- mcmc(att_sconserv.5)

mcmc_list <-
  mcmc.list(
    mcmc.att_sconserv.1,
    mcmc.att_sconserv.2,
    mcmc.att_sconserv.3,
    mcmc.att_sconserv.4,
    mcmc.att_sconserv.5
  )

is.mcmc.list(mcmc_list)
gelman.diag(mcmc_list)
gelman.plot(mcmc_list)


# relred

mcmc.att_relred.1 <- mcmc(att_relred.1)
mcmc.att_relred.2 <- mcmc(att_relred.2)
mcmc.att_relred.3 <- mcmc(att_relred.3)
mcmc.att_relred.4 <- mcmc(att_relred.4)
mcmc.att_relred.5 <- mcmc(att_relred.5)

mcmc_list <-
  mcmc.list(
    mcmc.att_relred.1,
    mcmc.att_relred.2,
    mcmc.att_relred.3,
    mcmc.att_relred.4,
    mcmc.att_relred.5
  )

is.mcmc.list(mcmc_list)
gelman.diag(mcmc_list)
gelman.plot(mcmc_list)

# ueskgen

mcmc.att_ueskgen.1 <- mcmc(att_ueskgen.1)
mcmc.att_ueskgen.2 <- mcmc(att_ueskgen.2)
mcmc.att_ueskgen.3 <- mcmc(att_ueskgen.3)
mcmc.att_ueskgen.4 <- mcmc(att_ueskgen.4)
mcmc.att_ueskgen.5 <- mcmc(att_ueskgen.5)

mcmc_list <-
  mcmc.list(
    mcmc.att_ueskgen.1,
    mcmc.att_ueskgen.2,
    mcmc.att_ueskgen.3,
    mcmc.att_ueskgen.4,
    mcmc.att_ueskgen.5
  )

is.mcmc.list(mcmc_list)
gelman.diag(mcmc_list)
gelman.plot(mcmc_list)

# govexp

mcmc.att_govexp.1 <- mcmc(att_govexp.1)
mcmc.att_govexp.2 <- mcmc(att_govexp.2)
mcmc.att_govexp.3 <- mcmc(att_govexp.3)
mcmc.att_govexp.4 <- mcmc(att_govexp.4)
mcmc.att_govexp.5 <- mcmc(att_govexp.5)

mcmc_list <-
  mcmc.list(
    mcmc.att_govexp.1,
    mcmc.att_govexp.2,
    mcmc.att_govexp.3,
    mcmc.att_govexp.4,
    mcmc.att_govexp.5
  )

is.mcmc.list(mcmc_list)
gelman.diag(mcmc_list)
gelman.plot(mcmc_list)

# END ==========================================================================



# Saving on object in RData format --------------------------------------------------------------------------------

#save(data.final.2, file = "raw_data_rep_1.2.23.RData")
#
#write_csv(data.final.2, "raw_data_rep_1.2.23.csv")
#
#save.image(file = "raw_workspace_1.2.23.RData")